package com.northpool.operator.statistics.dataset.math;

import com.alibaba.fastjson.JSON;
import com.northpool.operator.annotation.OperatorInfo;
import com.northpool.operator.common.algorithm.annotation.PluginInfo;
import com.northpool.operator.param.ParamTool;
import com.northpool.operator.statistics.dataset.helper.DataPointInfo;
import com.northpool.operator.type.TypeEnum;
import com.northpool.resources.command.QueryFilter;
import com.northpool.resources.datatable.FieldEncoder;
import com.northpool.resources.datatable.ITable;

import com.northpool.operator.annotation.InputParam;
import com.northpool.operator.annotation.InputParams;
import com.northpool.operator.annotation.OutputParam;
import com.northpool.operator.annotation.OutputParams;
import org.apache.commons.math3.ml.clustering.CentroidCluster;
import org.apache.commons.math3.ml.clustering.KMeansPlusPlusClusterer;

import java.util.*;
import java.util.stream.Collectors;

@OperatorInfo(
        ref = "StatisticalKMeansRegression",
        name = "KMeans聚类分析"
)
@OutputParams({
        @OutputParam(ref = "statisticalResults", name = "统计结果", type = TypeEnum.STRING)
})
public class StatisticalKMeansRegression extends BaseMathStatistical {

    @InputParams({
            @InputParam(ref = "dataSourceId", name = "数据源ID", desc = "数据源ID", type = TypeEnum.DATASOURCE, required = true, testValue = "402881ad9310bde101931141d66a0015"),
            @InputParam(ref = "tableName", name = "数据表名称", desc = "数据表名称，如: public.table", type = TypeEnum.STRING, testValue = "student_score_new"),
            @InputParam(ref = "infoFieldName", name = "字段名称", desc = "字段名称", type = TypeEnum.STRING, testValue = "userid"),
            @InputParam(ref = "valueFieldName", name = "值字段名称", desc = "值字段名称", type = TypeEnum.STRING, testValue = "math_score"),
            @InputParam(ref = "nClusters", name = "聚类分组数", desc = "聚类分组数", type = TypeEnum.NUMBER, testValue = "6")
    })
    public StatisticalKMeansRegression() {
    }

    @Override
    protected void executeStatistical(ParamTool paramTool) {
        String infoFieldName= paramTool.getStr("infoFieldName");
        String valueFieldName= paramTool.getStr("valueFieldName");
        List<String> listValueFields= Arrays.stream(valueFieldName.split(",")).collect(Collectors.toList());
        listValueFields.add(0,infoFieldName);

        try {

            String[] fields = listValueFields.stream().map(a->a).toArray(String[]::new);
            ITable table = tableOperator.getTable(fields, FieldEncoder.ORIGIN_FIELD_ENCODER);

            QueryFilter filter = new QueryFilter();
            List<Map<String, Object>> queryResultData = table.mapDao().query(filter);

            List<DataPointInfo> listDataPointInfo = new ArrayList<>();
            for (int i = 0; i < queryResultData.size(); i++) {
                Map<String, Object> stringObjectMap = queryResultData.get(i);
                String strInfoFieldValue= stringObjectMap.get(infoFieldName).toString();

                double[] valueArr=new double[listValueFields.size()];
                for (int fieldIndex = 0; fieldIndex < listValueFields.size(); fieldIndex++) {
                    valueArr[fieldIndex]= Double.parseDouble(stringObjectMap.get(listValueFields.get(fieldIndex)).toString());
                }
                listDataPointInfo.add(new DataPointInfo(strInfoFieldValue, valueArr));
            }

//            this.progress(99);

            // 设置聚类数量
            int nClusters = Integer.parseInt(paramTool.getStr("nClusters"));

            // K-means 聚类
            KMeansPlusPlusClusterer<DataPointInfo> clusterer = new KMeansPlusPlusClusterer<>(nClusters);
            List<CentroidCluster<DataPointInfo>> clusters = clusterer.cluster(listDataPointInfo);

            Map<String,List<DataPointInfo>> mapKMeansResult=new LinkedHashMap<>();
            // 输出每个学生的成绩及其聚类标签
            Integer index=0;
            for (CentroidCluster<DataPointInfo> cluster : clusters) {
                index++;
                List<DataPointInfo> list=new LinkedList<>();
                for (DataPointInfo dataPointInfo : cluster.getPoints()) {
                    list.add(dataPointInfo);
                }
                mapKMeansResult.put("K-Means "+index,list);
            }

            this.writeResult(JSON.toJSONString(mapKMeansResult));
            this.successExit();
        }catch (Exception ex) {
            this.error(ex.getMessage());
            this.failedExit();
        }
    }

    public static void main(String[] args) {
        // 创建示例成绩数据，假设有多个科目
        int numStudents = 50;
        int numSubjects = 2; // 假设有 3 个科目
        List<DataPointInfo> listDataPointInfo = new ArrayList<>();

        // 随机生成学生成绩
        for (int i = 0; i < numStudents; i++) {
            String studentId = "Student_" + (i + 1);
            double[] studentScores = new double[numSubjects];
            for (int j = 0; j < numSubjects; j++) {
                studentScores[j] = 50 + Math.random() * 50; // 随机成绩
            }
            listDataPointInfo.add(new DataPointInfo(studentId, studentScores));
        }

        // 设置聚类数量
        int nClusters = 3;

        // K-means 聚类
        KMeansPlusPlusClusterer<DataPointInfo> clusterer = new KMeansPlusPlusClusterer<>(nClusters);
        List<CentroidCluster<DataPointInfo>> clusters = clusterer.cluster(listDataPointInfo);

        // 输出每个学生的成绩及其聚类标签
        for (CentroidCluster<DataPointInfo> cluster : clusters) {
            System.out.println("Cluster:");
            for (DataPointInfo student : cluster.getPoints()) {
                System.out.printf("%s: Scores = ", student.getId());
                for (double score : student.getDataValues()) {
                    System.out.printf("%.2f ", score);
                }
                System.out.println();
            }
            System.out.println();
        }
    }
}